package _01_第一季.树04.Tree01;

import java.util.Comparator;

@SuppressWarnings("unchecked")
public class BST<E> extends BinaryTree<E>{
    private Comparator<E> comparator; //支持元素可比较性

    public BST() {
        this(null);
    }

    public BST(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    private void elementNotNullCheck(E element){
        if(element==null){
            throw new IllegalArgumentException("element can not be null");
        }
    }
    public boolean contains(E element){
        return node(element) != null;
    }

    private int compare(E e1,E e2){
        if(comparator!=null){
            return comparator.compare(e1, e2);
        }
        return ((Comparable<E>)e1).compareTo(e2);
    }

    //根据元素找到节点
    private Node<E> node(E element){
        elementNotNullCheck(element);
        Node<E> node = root;
        while(node!=null){
            int cmp = compare(element, node.element);
            if(cmp == 0) return node;
            if(cmp>0){
                node = node.right;
            }else{
                node = node.left;
            }
        }
        return null;
    }

    //添加元素
    public void add(E element){
        elementNotNullCheck(element);

        if(root == null){
            root = createNode(element,null);
            size++;
            afterAdd(root);
            return;
        }

        Node<E> node = root;
        Node<E> parent = root;
        int cmp = 0;
        while(node!=null){
            cmp = compare(element, node.element);
            parent = node;
            if(cmp>0){
                node = node.right;
            }else if(cmp<0){
                node = node.left;
            }else{
                node.element = element;
                return;
            }
        }

        Node<E> newNode = createNode(element,parent);
        if(cmp>0){
            newNode.parent.right = newNode;
        }else{
            newNode.parent.left = newNode;
        }
        size++;
        afterAdd(newNode);
    }

    //添加之后的调整
    protected void afterAdd(Node<E> node){

    }

    //删除之后的调整
    protected void afterRemove(Node<E> node){

    }

    public void remove(E element){
        remove(node(element));
    }

    private void remove(Node<E> node){
        if(node == null) return;
        size--;

        if(node.hasTwoKids()){
            Node<E> pre = predecessor(node);
            node.element = pre.element;
            node = pre;
        }

        Node<E> replacement = node.left!=null ? node.left : node.right;
        if(replacement!=null){
            replacement.parent = node.parent;
            if(node.parent == null){
                root = replacement;
            }else if(node.isLeftChild()){
                node.parent.left = replacement;
            }else{
                node.parent.right = replacement;
            }
            afterRemove(replacement);
        }else if(node.parent == null){
            root = null;
            afterRemove(node);
        }else{
            if(node.isLeftChild()){
                node.parent.left = null;
            }else{
                node.parent.right = null;
            }
            afterRemove(node);
        }
    }
}
